import gradio as gr
import scipy.io as scio
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties
import io
import sys
import os
import tempfile
import base64
from scipy.fftpack import fft
from sklearn.decomposition import PCA
from PIL import Image
import pandas as pd

# 获取当前文件的目录
current_dir = os.path.dirname(os.path.abspath(__file__))

# 添加上级目录的模块路径到sys.path
sys.path.append(os.path.join(current_dir, "model"))

from main import invoke_model

# 设置 Matplotlib 使用的字体
plt.rcParams["font.sans-serif"] = ["SimHei"]  # 使用黑体
plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示问题

# 定义全局变量
x_final = []


def display_images(file):
    global x_final  # 声明 x_final 为全局变量
    data = scio.loadmat(file.name)
    label = data["label"]
    y1 = data["data1"]
    phi1 = data["phi1"]

    y2 = data["data2"]
    phi2 = data["phi2"]

    label_array = np.array(label)

    y_array1 = np.array(y1)
    phi_array1 = np.array(phi1)

    y_array2 = np.array(y2)
    phi_array2 = np.array(phi2)

    x_data1 = []
    x_data2 = []
    x_data = []
    for i in range(len(label_array)):
        x_data1.append(
            np.dot(np.linalg.pinv(phi_array1[i]), y_array1[i])
        )  # np.linalg.pinv意思为计算伪逆
        x_data2.append(np.dot(np.linalg.pinv(phi_array2[i]), y_array2[i]))

    x_data1 = np.array(x_data1)
    x_data2 = np.array(x_data2)

    x_data1_real = np.real(x_data1)
    x_data1_imag = np.imag(x_data1)
    x_data2_real = np.real(x_data2)
    x_data2_imag = np.imag(x_data2)

    x_data1_real = x_data1_real.reshape(-1, 1)
    x_data2_real = x_data2_real.reshape(-1, 1)

    x_data1_imag = x_data1_imag.reshape(-1, 1)
    x_data2_imag = x_data2_imag.reshape(-1, 1)

    x_data3_real = np.concatenate((x_data1_real, x_data2_real), axis=1)
    x_data3_imag = np.concatenate((x_data1_imag, x_data2_imag), axis=1)

    # 选择一个切片进行可视化（假设是第一个样本的频谱数据）
    x_real = x_data3_real  # 实部
    x_imag = x_data3_imag  # 虚部

    # 定义时间轴（根据数据情况）
    L = 195
    R = 1
    K = 91
    K0 = 10
    fnyq = 10e10
    TimeResolution = 1 / fnyq
    TimeWin = [0, L * R * K - 1, L * R * (K + K0) - 1]
    t_axis = np.arange(TimeWin[0], TimeWin[-1] + 1) * TimeResolution

    # 定义数字时间轴
    Digital_time_axis = np.linspace(t_axis[0], t_axis[-1], x_real.shape[0])

    # 要可视化的样本数据（选择一个特定的频率分量，例如第一个）
    DigitalSamples1 = x_real
    DigitalSignalSamples = x_imag

    # 绘制 DigitalSamples1
    time_axis_min = 0
    time_axis_max = 2e-7

    # 振幅范围可以手动设置为较小的范围
    amplitude_min = -5000
    amplitude_max = 5000

    # 绘图
    plt.figure(figsize=(10, 6))  # 调整图像大小
    plt.plot(Digital_time_axis, DigitalSamples1, "r")
    plt.title("预处理前")
    plt.xlabel("times (s)")
    plt.ylabel("magnitude")
    plt.xlim(time_axis_min, time_axis_max)
    plt.ylim(amplitude_min, amplitude_max)
    plt.grid(True)
    plt.tight_layout(pad=3.0)  # 确保坐标轴标题不被截断
    plt.subplots_adjust(right=0.95)  # 调整右边距

    # 将图像保存为NumPy数组
    img_1 = io.BytesIO()
    plt.savefig(
        img_1, format="png", bbox_inches="tight"
    )  # 使用 bbox_inches='tight' 确保图像不被截断
    img_1.seek(0)
    img_1 = Image.open(img_1)
    plt.close()  # 关闭当前图像，防止内存泄漏

    pca = PCA(1)
    x_data_real = pca.fit_transform(x_data3_real)
    x_data_imag = pca.fit_transform(x_data3_imag)

    x_data_real = x_data_real.reshape(1 * 195 * 101, 1)  # 根据你产生的数据集修改大小
    x_data_imag = x_data_imag.reshape(1 * 195 * 101, 1)

    for j in range(len(x_data_imag)):
        x_data.append(complex(x_data_real[j, 0], x_data_imag[j, 0]))
    x_data = np.array(x_data)
    x_data = x_data.reshape(1, 195, 101)

    x_data_array = np.array(x_data)
    x_data_array_fft = fft(x_data_array)

    x_data_real = np.real(x_data_array_fft)
    x_data_img = np.imag(x_data_array_fft)

    x_final = np.stack((x_data_real, x_data_img), axis=3)

    # 选择一个切片进行可视化（假设是第一个样本的频谱数据）
    x_real = x_final[0, :, :, 0]  # 实部
    x_imag = x_final[0, :, :, 1]  # 虚部

    # 定义时间轴（根据数据情况）
    L = 195
    R = 1
    K = 91
    K0 = 10
    fnyq = 10e10
    TimeResolution = 1 / fnyq
    TimeWin = [0, L * R * K - 1, L * R * (K + K0) - 1]
    t_axis = np.arange(TimeWin[0], TimeWin[-1] + 1) * TimeResolution

    # 定义数字时间轴
    Digital_time_axis = np.linspace(t_axis[0], t_axis[-1], x_real.shape[0])

    # 要可视化的样本数据（选择一个特定的频率分量，例如第一个）
    DigitalSamples1 = x_real
    DigitalSignalSamples = x_imag

    # 绘制 DigitalSamples1
    time_axis_min = 0
    time_axis_max = 2e-7

    # 振幅范围可以手动设置为较小的范围
    amplitude_min = -30000
    amplitude_max = 30000

    # 绘图
    plt.figure(figsize=(10, 6))  # 调整图像大小
    plt.plot(Digital_time_axis, DigitalSamples1, "r")
    plt.title("预处理后")
    plt.xlabel("times (s)")
    plt.ylabel("magnitude")
    plt.xlim(time_axis_min, time_axis_max)
    plt.ylim(amplitude_min, amplitude_max)
    plt.grid(True)
    plt.tight_layout(pad=3.0)  # 确保坐标轴标题不被截断
    plt.subplots_adjust(right=0.95)  # 调整右边距

    # 将图像保存为NumPy数组
    img_2 = io.BytesIO()
    plt.savefig(
        img_2, format="png", bbox_inches="tight"
    )  # 使用 bbox_inches='tight' 确保图像不被截断
    img_2.seek(0)
    img_2 = Image.open(img_2)
    plt.close()  # 关闭当前图像，防止内存泄漏

    return img_1, img_2


def process_data(file):
    data = scio.loadmat(file.name)
    label = data["label"]
    global x_final  # 声明 x_final 为全局变量
    data_out = {"x": x_final, "label": label}
    temp_mat_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mat")
    scio.savemat(temp_mat_file.name, data_out)
    temp_mat_file.close()
    accuracy, output = invoke_model(temp_mat_file.name)
    if not temp_mat_file.closed:
        temp_mat_file.close()
    os.remove(temp_mat_file.name)

    # 将 Tensor 转换为 NumPy 数组
    output_np = output.numpy()

    # 将 NumPy 数组转换为整数类型
    output_np_int = output_np.astype(int)

    return accuracy, label, output_np_int


with gr.Blocks() as demo:
    file_input = gr.File(label="Upload .mat file")
    img1 = gr.Image()
    img2 = gr.Image()
    array1_output = gr.HTML()
    array2_output = gr.HTML()
    submit_btn = gr.Button("Submit")

    mean_output = gr.Textbox(label="Accuracy:")
    array1_output = gr.Textbox(label="输入信号的状态:")
    array2_output = gr.Textbox(label="模型的预判输出:")

    file_input.change(display_images, inputs=file_input, outputs=[img1, img2])
    submit_btn.click(
        process_data,
        inputs=file_input,
        outputs=[mean_output, array1_output, array2_output],
    )

demo.launch()
